use axum::{
    extract::{ws::{Message, WebSocketUpgrade}, State},
    response::IntoResponse
};
use futures_util::{SinkExt, StreamExt};

use std::sync::Arc;

use super::state::AppState;
// 服务器发送消息
pub async fn ws_handler( ws: WebSocketUpgrade, State(state): State<Arc<AppState>> ) -> impl IntoResponse {
    ws.on_upgrade(|socket| async move {
        let (mut sender, mut receiver) = socket.split();

        let mut username = String::new();
        // 循环直到找到一条文本消息
        while let Some(Ok(message)) = receiver.next().await {
            if let Message::Text(name) = message {
                // 如果未取客户端发送的 username，则填写 username 字符串
                if name == "ping" {
                   break;
                }
                check_username(&state, &mut username, &name);
                // 如果不为空，我们想退出循环，否则我们想退出函数
                if !username.is_empty() {
                    break;
                } else {
                    // 向我们的客户发送用户名已被占用的信息
                    let _ = sender.send(Message::Text(format!("{username} 名已被占用").into())).await;
                    return;
                }
            }
        }

        let mut rx = state.tx.subscribe();

        // 向所有订阅者发送已加入的消息
        _ = state.tx.send(format!("{username} 加入"));

        let mut send_task = tokio::spawn(async move {
            while let Ok(msg) = rx.recv().await {
                if sender.send(Message::Text(msg.into())).await.is_err() {
                    break;
                }
            }
        });

        let tx = state.tx.clone();
        
        let mut recv_task = tokio::spawn(async move {
            while let Some(Ok(Message::Text(text))) = receiver.next().await {
                if text == "ping" {
                    _ = tx.send(format!("pong"));
                }else{
                    _ = tx.send(text.to_string());
                }
            }
        });
        tokio::select! {
            _ = &mut send_task => recv_task.abort(),
            _ = &mut recv_task => send_task.abort(),
        };
        // 用户离开消息
        let _ = state.tx.send(format!("{username} left."));
        // 从 Map 中删除用户名，以便新客户端可以再次使用它
        state.user_set.lock().unwrap().remove(&username);
    })
}

fn check_username(state: &AppState, string: &mut String, name: &str) {
    let mut user_set = state.user_set.lock().unwrap();

    if !user_set.contains(name) {
        user_set.insert(name.to_owned());
        string.push_str(name);
    }
}